"""控制理论模块
包含可控性、可观性、李雅普诺夫方程等功能
"""
import numpy as np
from scipy.linalg import solve_lyapunov, solve_continuous_are, solve_discrete_are
from typing import Tuple, Optional

class ControlTheory:
    @staticmethod
    def controllability_matrix(A: np.ndarray, B: np.ndarray) -> np.ndarray:
        """
        构造可控性矩阵
        :param A: 系统矩阵
        :param B: 输入矩阵
        :return: 可控性矩阵 [B, AB, A²B, ..., A^(n-1)B]
        """
        n = A.shape[0]
        controllability = B.copy()
        
        A_power = A.copy()
        for i in range(1, n):
            controllability = np.column_stack([controllability, A_power @ B])
            A_power = A_power @ A
        
        return controllability
    
    @staticmethod
    def observability_matrix(A: np.ndarray, C: np.ndarray) -> np.ndarray:
        """
        构造可观性矩阵
        :param A: 系统矩阵
        :param C: 输出矩阵
        :return: 可观性矩阵 [C; CA; CA²; ...; CA^(n-1)]
        """
        n = A.shape[0]
        observability = C.copy()
        
        A_power = A.copy()
        for i in range(1, n):
            observability = np.row_stack([observability, C @ A_power])
            A_power = A_power @ A
        
        return observability
    
    @staticmethod
    def is_controllable(A: np.ndarray, B: np.ndarray, tol: float = 1e-10) -> bool:
        """
        判断系统是否可控
        :param A: 系统矩阵
        :param B: 输入矩阵
        :param tol: 容差
        :return: 是否可控
        """
        controllability = ControlTheory.controllability_matrix(A, B)
        return bool(np.linalg.matrix_rank(controllability, tol=tol) == A.shape[0])
    
    @staticmethod
    def is_observable(A: np.ndarray, C: np.ndarray, tol: float = 1e-10) -> bool:
        """
        判断系统是否可观
        :param A: 系统矩阵
        :param C: 输出矩阵
        :param tol: 容差
        :return: 是否可观
        """
        observability = ControlTheory.observability_matrix(A, C)
        return bool(np.linalg.matrix_rank(observability, tol=tol) == A.shape[0])
    
    @staticmethod
    def controllability_gramian(A: np.ndarray, B: np.ndarray) -> np.ndarray:
        """
        计算可控性Gramian矩阵
        :param A: 系统矩阵
        :param B: 输入矩阵
        :return: 可控性Gramian矩阵
        """
        # 求解李雅普诺夫方程: A*Wc + Wc*A^T + B*B^T = 0
        return solve_lyapunov(A, -B @ B.T)
    
    @staticmethod
    def observability_gramian(A: np.ndarray, C: np.ndarray) -> np.ndarray:
        """
        计算可观性Gramian矩阵
        :param A: 系统矩阵
        :param C: 输出矩阵
        :return: 可观性Gramian矩阵
        """
        # 求解李雅普诺夫方程: A^T*Wo + Wo*A + C^T*C = 0
        return solve_lyapunov(A.T, -C.T @ C)
    
    @staticmethod
    def lyapunov_equation(A: np.ndarray, Q: np.ndarray) -> np.ndarray:
        """
        求解连续李雅普诺夫方程 A*X + X*A^T + Q = 0
        :param A: 系数矩阵
        :param Q: 右端矩阵
        :return: 解矩阵X
        """
        return solve_lyapunov(A, -Q)
    
    @staticmethod
    def discrete_lyapunov_equation(A: np.ndarray, Q: np.ndarray) -> np.ndarray:
        """
        求解离散李雅普诺夫方程 A*X*A^T - X + Q = 0
        :param A: 系数矩阵
        :param Q: 右端矩阵
        :return: 解矩阵X
        """
        return solve_lyapunov(A, Q, discrete=True)
    
    @staticmethod
    def continuous_riccati_equation(A: np.ndarray, B: np.ndarray, 
                                   Q: np.ndarray, R: np.ndarray) -> np.ndarray:
        """
        求解连续代数Riccati方程 (CARE)
        A^T*X + X*A - X*B*R^(-1)*B^T*X + Q = 0
        :param A: 系统矩阵
        :param B: 输入矩阵
        :param Q: 状态权重矩阵
        :param R: 输入权重矩阵
        :return: 解矩阵X
        """
        return solve_continuous_are(A, B, Q, R)
    
    @staticmethod
    def discrete_riccati_equation(A: np.ndarray, B: np.ndarray,
                                 Q: np.ndarray, R: np.ndarray) -> np.ndarray:
        """
        求解离散代数Riccati方程 (DARE)
        A^T*X*A - X - A^T*X*B*(R + B^T*X*B)^(-1)*B^T*X*A + Q = 0
        :param A: 系统矩阵
        :param B: 输入矩阵
        :param Q: 状态权重矩阵
        :param R: 输入权重矩阵
        :return: 解矩阵X
        """
        return solve_discrete_are(A, B, Q, R)
    
    @staticmethod
    def lqr_controller(A: np.ndarray, B: np.ndarray, 
                      Q: np.ndarray, R: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        设计LQR控制器
        :param A: 系统矩阵
        :param B: 输入矩阵
        :param Q: 状态权重矩阵
        :param R: 输入权重矩阵
        :return: (反馈增益矩阵K, Riccati方程解P)
        """
        # 验证输入矩阵的性质
        if not ControlTheory._is_positive_semidefinite(Q):
            raise ValueError("Q matrix must be positive semidefinite")
        if not ControlTheory._is_positive_definite(R):
            raise ValueError("R matrix must be positive definite")
        
        # 检查系统可控性
        if not ControlTheory.is_controllable(A, B):
            import warnings
            warnings.warn("System is not completely controllable, LQR may not be optimal")
        
        try:
            P = solve_continuous_are(A, B, Q, R)
            # 使用数值稳定的方法计算反馈增益
            K = np.linalg.solve(R, B.T @ P)
            return K, P
        except Exception as e:
            raise ValueError(f"Failed to solve Riccati equation: {str(e)}")
    
    @staticmethod
    def _is_positive_definite(matrix: np.ndarray, tol: float = 1e-8) -> bool:
        """检查矩阵是否正定"""
        try:
            eigenvals = np.linalg.eigvals(matrix)
            return np.all(eigenvals > tol)
        except:
            return False
    
    @staticmethod
    def _is_positive_semidefinite(matrix: np.ndarray, tol: float = 1e-8) -> bool:
        """检查矩阵是否半正定"""
        try:
            eigenvals = np.linalg.eigvals(matrix)
            return np.all(eigenvals >= -tol)
        except:
            return False